import torch.utils.data as data_utils
import torch
import pickle


class MyDataset(data_utils.Dataset):
    def __init__(self, max_len, key='train_data') -> None:
        super().__init__()
        self.max_len = max_len
        with open('../../data.bin', 'rb+') as f:
            seq = pickle.load(f)
            self.item_seq = [ele[0] for ele in seq[key]]        # 历史商品交互序列
            self.query_seq = [ele[1] for ele in seq[key]]       # 历史查询词序列的序列
            self.cur_item = [ele[2] for ele in seq[key]]        # 当前交互商品id
            self.cur_query = [ele[3] for ele in seq[key]]       # 当前查询词id

    def __len__(self):
        return len(self.item_seq)

    def __getitem__(self, index):
        # return torch.LongTensor(self.item_seq[index]), torch.LongTensor([self.cur_item[index]])
        return torch.LongTensor(self.item_seq[index]), torch.LongTensor([self.cur_item[index]]), \
                    torch.LongTensor(self.query_seq[index]), torch.LongTensor([self.cur_query[index]])


class TrainDataset(data_utils.Dataset):
    def __init__(self, id2seq, max_len):
        self.id2seq = id2seq
        self.max_len = max_len

    def __len__(self):
        return len(self.id2seq)

    def __getitem__(self, index):
        seq = self._getseq(index)
        labels = [seq[-1]]
        tokens = seq[:-1]
        tokens = tokens[-self.max_len:]
        mask_len = self.max_len - len(tokens)
        tokens = [0] * mask_len + tokens
        return torch.LongTensor(tokens), torch.LongTensor(labels)

    def _getseq(self, idx):
        return self.id2seq[idx]


class Data_Train():
    def __init__(self, data_train, args):
        self.u2seq = data_train
        self.max_len = args.max_len
        self.batch_size = args.batch_size
        self.split_onebyone()

    def split_onebyone(self):
        self.id_seq = {}
        self.id_seq_user = {}
        idx = 0
        for user_temp, seq_temp in self.u2seq.items():
            for star in range(len(seq_temp)-1):
                self.id_seq[idx] = seq_temp[:star+2]
                self.id_seq_user[idx] = user_temp
                idx += 1

    def get_pytorch_dataloaders(self):
        # dataset = TrainDataset(self.id_seq, self.max_len)
        dataset = MyDataset(15, 'train_data')
        return data_utils.DataLoader(dataset, batch_size=self.batch_size, shuffle=True, pin_memory=True)


class ValDataset(data_utils.Dataset):
    def __init__(self, u2seq, u2answer, max_len):
        self.u2seq = u2seq  # data_train
        self.users = sorted(self.u2seq.keys())  # 0 - len(data_train) - 1
        self.u2answer = u2answer    # data_val
        self.max_len = max_len

    def __len__(self):
        return len(self.users)

    def __getitem__(self, index):
        user = self.users[index]    # 随便挑一个用户
        seq = self.u2seq[user]      # 取出用户对应的训练数据
        answer = self.u2answer[user]# 取出验证集用户对应的数据
        seq = seq[-self.max_len:]   # 截断训练数据
        padding_len = self.max_len - len(seq)
        seq = [0] * padding_len + seq
        return torch.LongTensor(seq),  torch.LongTensor(answer)


class Data_Val():
    def __init__(self, data_train, data_val, args):
        self.batch_size = args.batch_size
        self.u2seq = data_train
        self.u2answer = data_val
        self.max_len = args.max_len
        

    def get_pytorch_dataloaders(self):
        # dataset = ValDataset(self.u2seq, self.u2answer, self.max_len)
        dataset = MyDataset(15, 'eval_data')
        print(len(dataset))
        dataloader = data_utils.DataLoader(dataset, batch_size=self.batch_size, shuffle=False, pin_memory=True)
        return dataloader


class TestDataset(data_utils.Dataset):
    def __init__(self, u2seq, u2_seq_add, u2answer, max_len):
        self.u2seq = u2seq
        self.u2seq_add = u2_seq_add
        self.users = sorted(self.u2seq.keys())
        self.u2answer = u2answer
        self.max_len = max_len

    def __len__(self):
        return len(self.users)

    def __getitem__(self, index):
        user = self.users[index]
        seq = self.u2seq[user] + self.u2seq_add[user]
        # seq = self.u2seq[user]
        answer = self.u2answer[user]
        seq = seq[-self.max_len:]
        padding_len = self.max_len - len(seq)
        seq = [0] * padding_len + seq
        return torch.LongTensor(seq), torch.LongTensor(answer)


class Data_Test():
    def __init__(self, data_train, data_val, data_test, args):
        self.batch_size = args.batch_size
        self.u2seq = data_train
        self.u2seq_add = data_val
        self.u2answer = data_test
        self.max_len = args.max_len

    def get_pytorch_dataloaders(self):
        # dataset = TestDataset(self.u2seq, self.u2seq_add, self.u2answer, self.max_len)
        dataset = MyDataset(15, 'test_data')
        print(len(dataset))
        dataloader = data_utils.DataLoader(dataset, batch_size=self.batch_size, shuffle=False, pin_memory=True)
        return dataloader


class CHLSDataset(data_utils.Dataset):
    def __init__(self, data, max_len):
        self.data = data
        self.max_len = max_len

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):

        data_temp = self.data[index]
        seq = data_temp[:-1]
        answer = [data_temp[-1]]
        seq = seq[-self.max_len:]
        padding_len = self.max_len - len(seq)
        seq = [0] * padding_len + seq
        return torch.LongTensor(seq), torch.LongTensor(answer)


class Data_CHLS():
    def __init__(self, data, args):
        self.batch_size = args.batch_size
        self.max_len = args.max_len
        self.data = data

    def get_pytorch_dataloaders(self):
        dataset = CHLSDataset(self.data, self.max_len)
        dataloader = data_utils.DataLoader(dataset, batch_size=self.batch_size, shuffle=False, pin_memory=True)
        return dataloader
